from zhipuai import ZhipuAI
from post_process import sql2bql
import argparse
import re
import json
import os
import logging

logger = logging.getLogger('nl2bql_logger')

# 加载配置
def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.json')
    with open(config_path, 'r') as file:
        config = json.load(file)
    
    return config
# 加载业务实体压缩文件
def load_database_schema():
    config=load_config()
    be_path=config['be_path']
    with open(be_path, 'r', encoding='utf-8') as file:
        database_schema = file.read()
    return database_schema



base_instruction="""
### Task

生成一个SQL来回答如下问题:

{question}

### Database Schema

查询语句基于如下使用字符串描述的数据库:
{database_schema}
"""



# 获取ZhipuAI的Completion结果
def get_zhipu_completion(question):
    config=load_config()
    client = ZhipuAI(api_key=config['zhipu']['api_key']) 
    database_scheme=load_database_schema()
    prompt=base_instruction.format(question=question, database_schema=database_scheme)
    response = client.chat.completions.create(
        model="codegeex-4",
        messages=[
        {
            "role": "system",
            "content": "你是一位SQL语言专家，你能帮助用于生成SQL语句或类似SQL语法规则的语句，并且能根据数据库的表结构给出SQL语句。你会为用户回答关于编程、代码、计算机方面的任何问题，并提供格式规范、可以执行、准确安全的代码，并在必要时提供详细的解释。"
        },
        {
            "role": "user",
            "content": prompt
        }
        ],
        top_p=0.7,
        temperature=0.9,
        max_tokens=1024,
        stop=["<|endoftext|>", "<|user|>", "<|assistant|>", "<|observation|>"]
    )
    return response.choices[0].message.content
# 提取代码块
def extract_code_blocks(text):
    pattern = r"```sql(.*?)```"
    code_blocks = re.findall(pattern, text,re.DOTALL)[0]
    return code_blocks


def main(input):
    # content='为了查询共有多少个销售组织，我们需要从`SalesOrganization`表中获取销售组织的数量。根据提供的数据库模式，`SalesOrganization`表的主键是`ID`， 所以我们可以通过以下SQL语句来查询销售组织的数量：\n\n```sql\nSELECT COUNT(*) AS TotalSalesOrganizations\nFROM SalesOrganization;\n```\n\n这条SQL语句的作用是计算`SalesOrganization`表中记录的总数，并将结果命名为`TotalSalesOrganizations`。'
    content=get_zhipu_completion(input)
    content=extract_code_blocks(content)
    content=sql2bql(content)
    return content
    

    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default='查询共有多少个销售组织', help="")
    
    args = parser.parse_args()
    main(args.input)
    